#include "main/utils.h"
using namespace std;

int getFirstK(const vector<int> &vec, const int &k) {
  int begin = 0;
  int end = vec.size() - 1;
  while (begin <= end) {
    int mid = begin + ((end - begin) >> 1);
    if (vec[mid] < k) {
      begin = mid + 1;
    } else if (vec[mid] > k) {
      end = mid - 1;
    } else if (mid == 0 || vec[mid - 1] != k) {
      return mid;
    } else {
      end = mid - 1;
    }
  }
  return -1;
}

int getLastK(const vector<int> &vec, const int &k) {
  int begin = 0;
  int end = vec.size() - 1;
  while (begin <= end) {
    int mid = begin + ((end - begin) >> 1);
    if (vec[mid] < k) {
      begin = mid + 1;
    } else if (vec[mid] > k) {
      end = mid - 1;
    } else if (mid == vec.size() - 1 || vec[mid + 1] != k) {
      return mid;
    } else {
      begin = mid + 1;
    }
  }
  return -1;
}

int getNumOfK(vector<int> vec, int k) {
  int idx_begin = getFirstK(vec, k);
  int idx_end = getLastK(vec, k);
  return idx_end - idx_begin + 1;
}

int getMissingNum(const vector<int> &vec) {
  int begin = 0;
  int end = vec.size() - 1;
  while (begin <= end) {
    int mid = begin + ((end - begin) >> 1);
    if (vec[mid] != mid && ((mid == 0) || vec[mid - 1] == mid - 1)) {
      return mid;
    }
    if (vec[mid] == mid) {
      begin = mid + 1;
    } else {
      end = mid - 1;
    }
  }
  return -1;
}

int equalNumAndIndex(const vector<int> &vec) {
  int begin = 0;
  int end = vec.size() - 1;
  while (begin <= end) {
    int mid = begin + ((end - begin) >> 1);
    if (vec[mid] == mid) {
      return mid;
    }
    if (vec[mid] < mid) {
      begin = mid + 1;
    } else {
      end = mid - 1;
    }
  }
  return -1;
}

int main() {
  vector<int> vec = {3, 3, 3, 3, 4, 5, 6, 7, 8};
  int num = getNumOfK(vec, 3);
  cout << "3 occurs " << num << " times" << endl;

  vec = {0, 1, 2, 3, 4, 6, 7, 8, 9, 10};
  num = getMissingNum(vec);
  cout << "The missing number is: " << num << endl;

  vec = {-3, -1, 1, 3, 5};
  num = equalNumAndIndex(vec);
  cout << "The element of number equal with index: " << num << endl;

  return 0;
}
